#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void declare_global(CodeGen *codegen); // 声明全局变量 a
void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    declare_global(codegen);

    codegen->append_inst(".text", ASMInstruction::Atrribute);
    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

void declare_global(CodeGen *codegen) {
    // 声明全局变量 a
    // 将 a 放到 .data 段中
    codegen->append_inst(".data",ASMInstruction::Atrribute);
    // 标记 a 为全局变量
    codegen->append_inst(".globl a", ASMInstruction::Atrribute);
    // 标记 a 为数据对象 (变量)
    codegen->append_inst(".type a, @object", ASMInstruction::Atrribute);
    // 标记 a 的大小为 4 字节
    codegen->append_inst(".size a, 4", ASMInstruction::Atrribute);
    // a 的标签
    codegen->append_inst("a", ASMInstruction::Label);
    // 分配 4 字节未初始化空间
    codegen->append_inst(".zero 4");
}

// TODO: 按照自己的理解翻译main函数，可以按照注释中的步骤，也可自由发挥
// 提示：可以使用 la t0, a  指令将全局变量a的内存地址放到t0寄存器中
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明 main 函数 */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi sp, sp, -32");
    codegen->append_inst("sd ra, 24(sp)");
    codegen->append_inst("sd s0, 16(sp)");
    codegen->append_inst("addi s0, sp, 32");

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);

    /* store i32 10, i32* @a */
    codegen->append_inst("store i32 10, i32* @a", ASMInstruction::Comment);
    codegen->append_inst("la", {"t0", "a"});
    codegen->append_inst("li", {"t1", "10"});
    codegen->append_inst("sw", {"t1", "0(t0)"});

    /* %op0 = load i32, i32* @a */
    codegen->append_inst("%op0 = load i32, i32* @a", ASMInstruction::Comment);
    offset_map["%op0"] = -24;
    codegen->append_inst("la", {"t0", "a"});
    codegen->append_inst("lw", {"t1", "0(t0)"});
    codegen->append_inst("sw", {"t1", "s0", "-24"});

    /* ret i32 %op0 */
    codegen->append_inst("ret i32 %op0", ASMInstruction::Comment);
    codegen->append_inst("lw", {"a0", "s0", "-24"});
    codegen->append_inst("j main_exit");

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 24(sp)");
    codegen->append_inst("ld s0, 16(sp)");
    codegen->append_inst("addi sp, sp, 32");
    codegen->append_inst("jr ra");
}